# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import abc
import datetime
import logging

import tabulate
from tornado.ioloop import IOLoop

import coin.flow.subscriber as flow_subscriber
from coin.base.param_util import to_list
from coin.flow.simple_queue import SimpleQueue


class PairStrategyBase(metaclass=abc.ABCMeta):
  ProductType1 = None
  ProductType2 = None
  LiveEnvironmentType1 = None
  LiveEnvironmentType2 = None
  OGName1 = "Unknown1"
  OGName2 = "Unknown1"

  def __init__(self, product_1, product_2, balance_currency=None, logger=None):
    self._logger = logger or logging.getLogger('PairStrategyBase')

    assert type(product_1) == self.ProductType1, (type(product_1), self.ProductType1)
    assert type(product_2) == self.ProductType2, (type(product_2), self.ProductType2)
    self._product_1 = product_1
    self._product_2 = product_2

    if balance_currency is None:
      self._balance_currency = balance_currency
    else:
      self._balance_currency = sorted(to_list(balance_currency))

    self._ioloop = None
    self._queue = None
    self._flow_sub = None

    self._og_1 = None
    self._og_2 = None

    self._book_1 = None
    self._book_2 = None

  def on_base_init(self, ioloop, queue, flow_sub):
    self._ioloop = ioloop
    self._queue = queue
    self._flow_sub = flow_sub

  def on_og_reset(self, name, og):
    self._logger.info('OG initialized: %s', name)
    if name == self.OGName1:
      self._og_1 = og
    elif name == self.OGName2:
      self._og_2 = og
    else:
      raise ValueError('Unknown OG: %s' % name)

    if self._og_2 is not None and self._og_1 is not None:
      self._start_print_balance()
      self._ioloop.add_callback(self.init_executor)

  def on_book_reset(self, name, book_builder):
    self._logger.info('Feed initialized: %s', name)
    if name == self.OGName1:
      book_builder.subscribe(self._product_1, self._on_book_1)
    elif name == self.OGName2:
      book_builder.subscribe(self._product_2, self._on_book_2)
    else:
      raise ValueError('Unknown feed: %s' % name)

  def is_og_ready(self):
    return (self._og_2 is not None and self._og_2.is_ready() and self._og_1 is not None
            and self._og_1.is_ready())

  def _print_balance(self):
    balance_1 = self._og_1.get_balances()
    balance_2 = self._og_2.get_balances()

    if self._balance_currency is not None:
      balance_currency = self._balance_currency
    else:
      balance_currency = sorted(set(balance_1.keys() + balance_2.keys()))

    table = []
    for cur in balance_currency:
      bal_1 = balance_1.get(cur, 0.)
      bal_2 = balance_2.get(cur, 0.)
      table.append(['*', cur, bal_1 + bal_2, bal_1, bal_2])

    self._logger.info(
        '----- Balance:\n* %s\n%s' %
        (str(datetime.datetime.now()),
         tabulate.tabulate(table,
                           headers=['*', 'Currency', 'Total', self.OGName1, self.OGName2],
                           floatfmt=".6f")))

  def _start_print_balance(self):
    try:
      if not self.is_og_ready():
        return
      self._print_balance()
    finally:
      self._ioloop.add_timeout(datetime.timedelta(seconds=30), self._start_print_balance)

  def _on_book_1(self, book):
    self._book_1 = book
    self.on_book()

  def _on_book_2(self, book):
    self._book_2 = book
    self.on_book()

  @abc.abstractmethod
  def init_executor(self):
    raise NotImplementedError()

  @abc.abstractmethod
  def on_book(self):
    raise NotImplementedError()

  @classmethod
  def run_live(cls,
               first_currency,
               first_product,
               first_og_config_name,
               second_product,
               second_og_config_name,
               on_base_init_callback,
               on_book_reset_callback,
               on_order_gateway_reset_callback,
               init_order_gateway_after=None,
               exit_after=None):
    # Initialize product
    first_product = to_list(first_product)
    for p in first_product:
      assert isinstance(p, cls.ProductType1)

    second_product = to_list(second_product)
    for p in second_product:
      assert isinstance(p, cls.ProductType2)

    # Initialize loop
    ioloop = IOLoop.current()
    queue = SimpleQueue('queue')
    flow_sub = flow_subscriber.from_queue(queue)
    on_base_init_callback(ioloop, queue, flow_sub)

    # Initialize first
    first_env = cls.LiveEnvironmentType1(first_currency,
                                         first_product,
                                         first_product,
                                         first_og_config_name,
                                         init_og_after=init_order_gateway_after,
                                         ioloop=ioloop,
                                         queue=queue,
                                         flow_sub=flow_sub)
    first_env.on_book_reset_callback.append(on_book_reset_callback)
    first_env.on_order_gateway_reset_callback.append(on_order_gateway_reset_callback)

    # Initialize second
    second_env = cls.LiveEnvironmentType2(second_product,
                                          second_product,
                                          second_og_config_name,
                                          init_og_after=init_order_gateway_after,
                                          ioloop=ioloop,
                                          queue=queue,
                                          flow_sub=flow_sub)
    second_env.on_book_reset_callback.append(on_book_reset_callback)
    second_env.on_order_gateway_reset_callback.append(on_order_gateway_reset_callback)

    first_env.start()
    second_env.start()

    if exit_after is not None:
      ioloop.add_timeout(exit_after, ioloop.stop)

    try:
      ioloop.start()
      return 0
    except KeyboardInterrupt:
      return 1
